"""
文本主题分类: ClueAI/PromptCLUE，只开放了PromptCLUE-base版本
"""

from threading import RLock
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks
from modelscope.models.nlp import T5ForConditionalGeneration
from modelscope.preprocessors import TextGenerationT5Preprocessor


class TextClassifier(object):

    lock = RLock()
    def __new__(cls, *args, **kwargs):
        with TextClassifier.lock:
            if not hasattr(TextClassifier, "_instance"):
                TextClassifier._instance = object.__new__(cls)
                repo_id = 'root/autodl-tmp/PromptCLUE'
                model = T5ForConditionalGeneration.from_pretrained(repo_id, revision='v0.1')
                preprocessor = TextGenerationT5Preprocessor(model.model_dir)
                cls.pipeline_t2t = pipeline(task=Tasks.text2text_generation, model=model, preprocessor=preprocessor)

        return TextClassifier._instance

    def __init__(self):
        self.class_names = ['故事', '文化', '娱乐', '体育', '财经', '房产', '汽车', '教育', '科技', '军事', '旅游', '国际', '股票', '农业', '游戏']

    def news_prompt(self, text: str) -> str:
        prompt = f'''这是关于哪方面的新闻：{text}
        选项：{','.join(self.class_names)}'''
        return prompt

    def predict(self, text: str) -> str:
        prompt = self.news_prompt(text)
        result = self.pipeline_t2t(prompt)
        return result.get('text', '')


if __name__ == '__main__':
    classifier = TextClassifier()
    print(classifier.predict("如果日本沉没，中国会接收日本难民吗？"))
